from typing import Optional, Callable

import torch.nn as nn
import torch
import torchvision.models.resnet as Resnet
from torch import Tensor

def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
    """3x3 convolution with padding"""
    return nn.Conv2d(
        in_planes,
        out_planes,
        kernel_size=3,
        stride=stride,
        padding=dilation,
        groups=groups,
        bias=False,
        dilation=dilation,
    )
class BasicBlock(nn.Module):
    expansion: int = 1

    def __init__(
        self,
        inplanes: int,
        planes: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        groups: int = 1,
        base_width: int = 64,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None,
    ) -> None:
        super().__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError("BasicBlock only supports groups=1 and base_width=64")
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)
            print("iden shape:",identity.shape)

        print("out shape:", out.shape)

        out += identity
        out = self.relu(out)

        return out

class MyBasicBlock(nn.Module):
    def __init__(self,in_channels,out_channels,downsample=None):
        super().__init__()
        self.conv1=nn.Conv2d(in_channels,out_channels,kernel_size=3,stride=1,padding=1, bias=False)
        self.bn1=nn.BatchNorm2d(out_channels)
        self.relu1=nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1,padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu2 = nn.ReLU(inplace=True)
        self.downsample=downsample
        if downsample is not None:
            self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2,padding=1, bias=False)
            self.downsample_func=nn.Conv2d(in_channels,out_channels,kernel_size=1,stride=2)
            self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self,x):
        indentity=x

        x=self.conv1(x)
        x=self.bn1(x)
        x=self.relu1(x)

        x = self.conv2(x)
        x = self.bn2(x)


        # print("x shape:",x.shape)
        # print("ind shape:", indentity.shape)

        if self.downsample is not None:
            indentity=self.downsample_func(indentity)
            indentity=self.bn2(indentity)

        # print("x shape:", x.shape)
        # print("ind shape:", indentity.shape)

        x= x+indentity
        x = self.relu2(x)

        return x

class MyBottleNeck(nn.Module):
    def __init__(self,in_channels,mid_channels, out_channels,downsample):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(mid_channels)
        self.relu1 = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(mid_channels)
        self.relu2 = nn.ReLU(inplace=True)

        self.conv3 = nn.Conv2d(mid_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels)
        self.relu3 = nn.ReLU(inplace=True)

        self.downsample = downsample
        if downsample is not None:
            self.conv2 = nn.Conv2d(mid_channels, mid_channels, kernel_size=3, stride=2, padding=1, bias=False)
            self.downsample_func = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=2)

    def forward(self,x):
        indentity=x

        x=self.conv1(x)
        x=self.bn1(x)
        x=self.relu1(x)

        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)

        x = self.conv3(x)
        x = self.bn3(x)
        x = self.relu3(x)

        # print("x shape:",x.shape)
        # print("ind shape:", indentity.shape)

        if self.downsample is not None:
            indentity=self.downsample_func(indentity)
            indentity=self.bn3(indentity)

        print("x shape:", x.shape)
        print("ind shape:", indentity.shape)

        x= x+indentity
        x = self.relu2(x)

        return x

# class resnet(nn.Module):
#     def __init__(self,in_channel,num_class,trans=None):
#         super(resnet,self).__init__()




if __name__=="__main__":
    # resnet34=Resnet.resnet34(pretrained=True)
    # print(resnet34)

    # input = torch.randn((16,64,56,56))
    # print("input shape:",input.shape)
    # downsample = nn.Sequential(
    #     conv1x1(64, 128, 2),
    #     nn.BatchNorm2d(128),
    # )
    # bb1=BasicBlock(64,128,stride=2, downsample=downsample)
    # output_1=bb1(input)
    # print("output1 shape:",output_1.shape)
    # # print(bb1)
    #
    # bb2=MyBasicBlock(64,128,downsample=True)
    # output_2=bb2(input)
    # print("output2 shape:",output_2.shape)
    # # print(bb2)

    input = torch.randn((16, 256, 56, 56))
    print("input shape:", input.shape)
    bn2=MyBottleNeck(256,128,512,downsample=True)
    output_2=bn2(input)
    print("output2 shape:", output_2.shape)


